package dev.alm.langchain4jspringbootdemo.config;

import dev.alm.langchain4jspringbootdemo.service.ToolService;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.service.*;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class AiConfig {

    public interface AiAssistant {
        String chat(String question);

        // 流式输出
        TokenStream stream(String question);
    }

    @Bean
    public AiAssistant assistant(ChatLanguageModel qwenChatModel, StreamingChatLanguageModel streamingQwenChatModel) {
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);
        return AiServices.builder(AiAssistant.class).chatLanguageModel(qwenChatModel).streamingChatLanguageModel(streamingQwenChatModel).chatMemory(chatMemory).build();
    }

    public interface ChatService {
        String chat(@MemoryId String memoryId, @UserMessage String message);

        @SystemMessage("你是一个票务助手。当用户发起有关退票要求时，必须始终要求用户提供姓名、班次以及退票原因。如若用户未提供，则直接返回让用户提供后再做处理。")
        TokenStream stream(@MemoryId String memoryId, @UserMessage String message);
    }

    public interface RagService {
        String chat(@MemoryId String memoryId, @UserMessage String message);
        TokenStream stream(@MemoryId String memoryId, @UserMessage String message);
    }

    @Bean
    public ChatService chatService(ChatLanguageModel qwenChatModel, StreamingChatLanguageModel streamingQwenChatModel, ToolService toolService) {
        return AiServices.builder(ChatService.class).tools(toolService).chatLanguageModel(qwenChatModel).streamingChatLanguageModel(streamingQwenChatModel).chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder().maxMessages(10).id(memoryId).build()).build();
    }

    public interface ChatServiceIdStore {
        String chat(@MemoryId String memoryId, @UserMessage String message);

        TokenStream stream(@MemoryId String memoryId, @UserMessage String message);
    }

    @Bean
    public ChatServiceIdStore chatServiceIdStore(ChatLanguageModel qwenChatModel, StreamingChatLanguageModel streamingQwenChatModel) {
        ChatStore store = new ChatStore();
        ChatMemoryProvider chatMemoryProvider = memoryId -> MessageWindowChatMemory.builder().maxMessages(10).id(memoryId).chatMemoryStore(store).build();
        return AiServices.builder(ChatServiceIdStore.class).chatLanguageModel(qwenChatModel).streamingChatLanguageModel(streamingQwenChatModel).chatMemoryProvider(chatMemoryProvider).build();
    }

}
